package cn.edu.thu.laud.jdbc.util;

import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.nio.charset.CharacterCodingException;
import java.util.zip.DataFormatException;
import java.util.zip.Deflater;
import java.util.zip.Inflater;

import org.apache.cassandra.io.util.FastByteArrayOutputStream;
import org.apache.cassandra.thrift.Compression;
import org.apache.cassandra.thrift.InvalidRequestException;
import org.apache.cassandra.utils.ByteBufferUtil;

import com.google.common.base.Charsets;

public class LaUDJdbcServerUtils {
	public static final short TYPE_CASSANDRA=1;
	public static final short TYPE_HIVE=2;
	public static final short TYPE_OTHER=3;
	public static final short TYPE_FILE=4;
	public static final char WRITE_INTO='w';
	public static final char READ_FROM='r';
	public static final char OTHERS='O';
	public static final short TYPE_MIX = 5;	
	   /**
     * Use the Compression object method to deflate the query string
     *
     * @param queryStr An un-compressed CQL query string
     * @param compression The compression object
     * @return A compressed string
     */
    public static ByteBuffer compressQuery(String queryStr, Compression compression)
    {
        byte[] data = queryStr.getBytes(Charsets.UTF_8);
        Deflater compressor = new Deflater();
        compressor.setInput(data);
        compressor.finish();

        ByteArrayOutputStream byteArray = new ByteArrayOutputStream();
        byte[] buffer = new byte[1024];

        try
        {
            while (!compressor.finished())
            {
                int size = compressor.deflate(buffer);
                byteArray.write(buffer, 0, size);
            }
        }
        finally
        {
            compressor.end(); //clean up after the Deflater
        }

       
        return ByteBuffer.wrap(byteArray.toByteArray());
    }
	/**
	 *  copy from CassandraServer.java .uncompress lasql
	 * @param query
	 * @param compression
	 * @return
	 * @throws InvalidRequestException
	 */
	public static String uncompress(ByteBuffer query, Compression compression) throws InvalidRequestException
	    {
	        String queryString = null;

	        // Decompress the query string.
	        try
	        {
	            switch (compression)
	            {
	                case GZIP:
	                    FastByteArrayOutputStream byteArray = new FastByteArrayOutputStream();
	                    byte[] outBuffer = new byte[1024], inBuffer = new byte[1024];

	                    Inflater decompressor = new Inflater();

	                    int lenRead = 0;
	                    while (true)
	                    {
	                        if (decompressor.needsInput())
	                            lenRead = query.remaining() < 1024 ? query.remaining() : 1024;
	                            query.get(inBuffer, 0, lenRead);
	                            decompressor.setInput(inBuffer, 0, lenRead);

	                        int lenWrite = 0;
	                        while ((lenWrite = decompressor.inflate(outBuffer)) !=0)
	                            byteArray.write(outBuffer, 0, lenWrite);

	                        if (decompressor.finished())
	                            break;
	                    }

	                    decompressor.end();

	                    queryString = new String(byteArray.toByteArray(), 0, byteArray.size(), "UTF-8");
	                    break;
	                case NONE:
	                    try
	                    {
	                        queryString = ByteBufferUtil.string(query);
	                    }
	                    catch (CharacterCodingException ex)
	                    {
	                        throw new InvalidRequestException(ex.getMessage());
	                    }
	                    break;
	            }
	        }
	        catch (DataFormatException e)
	        {
	            throw new InvalidRequestException("Error deflating query string.");
	        }
	        catch (UnsupportedEncodingException e)
	        {
	            throw new InvalidRequestException("Unknown query string encoding.");
	        }
	        return queryString;
	    }
}
